package com.zoffy.oj_backend.utils;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Consumer;

@Slf4j
@Component
public class SSEUtils {
    private final static Map<String, SseEmitter> sseEmitters = new HashMap<>();

    /**
     * 注册SSE连接
     *
     * @param userId 用户id
     * @return SSEEmitter
     */
    public static SseEmitter subscribe(String userId) {
        SseEmitter sseEmitter = new SseEmitter(30 * 1000L);
        sseEmitter.onTimeout(timeoutCallBack(userId));
        sseEmitter.onError(errorCallBack(userId));
        sseEmitters.put(userId, sseEmitter);
        return sseEmitter;
    }

    /**
     * 向指定SSE连接发送消息
     *
     * @param userId  用户id
     * @param content 消息文本
     */
    public static void push(String userId, String content) {
        SseEmitter sseEmitter = sseEmitters.get(userId);
        if (sseEmitter == null) {
            return;
        }
        try {
            sseEmitter.send(content);
        } catch (IOException e) {
            log.error("用户 {} SSE消息推送失败:{}", userId, e.getMessage());
            throw new RuntimeException(e);
        }
    }

    /**
     * 删除指定SSE连接
     *
     * @param userId 用户id
     */
    public static void removeSseEmitter(String userId) {
        SseEmitter sseEmitter = sseEmitters.get(userId);
        if (sseEmitter == null) {
            return;
        }
        sseEmitter.complete();
        sseEmitters.remove(userId);
    }

    private static Runnable timeoutCallBack(String userId) {
        return () -> removeSseEmitter(userId);
    }

    private static Consumer<Throwable> errorCallBack(String userId) {
        return throwable -> removeSseEmitter(userId);
    }
}
